import pandas as pd
from skimage.metrics import mean_squared_error
from sklearn.metrics import root_mean_squared_error, mean_absolute_error
from sklearn.preprocessing import StandardScaler, MinMaxScaler
from sklearn.preprocessing import StandardScaler
from xgboost import XGBRegressor
import xgboost as xgb
from sklearn.model_selection import train_test_split
from sklearn.metrics import roc_auc_score, roc_curve

def feature_engineering(data):
    stander = StandardScaler()
    # 复制源数据
    feature_date = data.copy()

    # 热编码
    object_data = pd.get_dummies(
        feature_date[['BusinessTravel', 'Department', 'JobRole',
                      'MaritalStatus', 'OverTime']])
    # object_data.drop(columns='Gender_Male', inplace=True)
    object_data.rename(columns={'Gender_Female': 'Gender'}, inplace=True)

    # 数值型特征
    numeric_features = ['JobInvolvement', 'JobLevel', 'MonthlyIncome', 'Age', 'EnvironmentSatisfaction',
                        'YearsInCurrentRole']
    # 归一化
    min_max_scaler = MinMaxScaler()
    normalized_features = min_max_scaler.fit_transform(feature_date[numeric_features])
    normalized_df = pd.DataFrame(normalized_features, columns=numeric_features)
    # print(normalized_df.info())

    # 拼接数据
    new_data = pd.concat([object_data, normalized_df], axis=1)
    print(new_data.info())
    return new_data, new_data.columns


def model_train(data_resource, data, columns):
    x = data[columns]
    y = data_resource['Attrition']
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=18)

    # 3.模型实例化
    es = xgb.XGBClassifier(n_estimators=150, max_depth=5, learning_rate=0.1)
    es.fit(x_train, y_train)
    y_pred = es.predict(x_test)
    y_pred_proba = es.predict_proba(x_test)
    y_score = y_pred_proba[:, 1]

    # 4、模型评估
    print(f"均方误差:{mean_squared_error(y_test, y_pred)}")
    print(f"均方根误差:{root_mean_squared_error(y_test, y_pred)}")
    print(f"平均绝对误差:{mean_absolute_error(y_test, y_pred)}")

    auc_value = roc_auc_score(y_test, y_score)
    print(f"The AUC score is: {auc_value:.4f}")

if __name__ == '__main__':
    pd_data = pd.read_csv('../../03_数据集/train.csv')
    feature_data, feature_columns = feature_engineering(pd_data)
    model_train(pd_data, feature_data, feature_columns)
